import redis
import json
import time
import copy
import numpy
from loguru import logger
from typing import List


class RedisClient:
    def __init__(self, host, port, db, max_connections):
        self.pool = redis.ConnectionPool(host=host, port=port, db=db, max_connections=max_connections)
        self.client = redis.Redis(connection_pool=self.pool)

    def get(self, key):
        try:
            return self.client.get(key)
        except redis.RedisError as e:
            # 处理Redis错误
            logger.error(f"Redis get error: {e}")
            return None

    def set(self, key, value):
        try:
            self.client.set(key, value)
        except redis.RedisError as e:
            # 处理Redis错误
            logger.error(f"Redis set error: {e}")

    def delete(self, key):
        try:
            self.client.delete(key)
        except redis.RedisError as e:
            # 处理Redis错误
            logger.error(f"Redis delete error: {e}")


class ChatHistory:
    """
    A class used to represent the chat history of a user.

    ...

    Attributes
    ----------
    uid : str
        a string representing the unique identifier of the user
    redis_client : RedisClient
        an instance of RedisClient used to interact with Redis
    history_len : int
        an integer representing the length of the history to be retrieved

    Methods
    -------
    get_history():
        Returns the chat history of the user from Redis.
    add_to_history(message: List):
        Adds a new message to the user's chat history in Redis.
    clear_history():
        Clears the user's chat history from Redis.
    """

    def __init__(self, uid, redis_client, history_len):
        """
        Constructs all the necessary attributes for the ChatHistory object.

        Parameters
        ----------
            uid : str
                a string representing the unique identifier of the user
            redis_client : RedisClient
                an instance of RedisClient used to interact with Redis
            history_len : int
                an integer representing the length of the history to be retrieved
        """
        self.uid = uid
        self.redis_client = redis_client
        self.history_len = history_len

    def get_history(self):
        """
        Retrieves the chat history of the user from Redis.

        The method retrieves the chat history, filters out any invalid dialogues
        (dialogues where the assistant's reply is empty), and returns the filtered history.

        Returns
        -------
        list
            a list of dialogues, each dialogue is a list of messages
        """
        try:
            # 从Redis中获取历史对话记录
            history_json = self.redis_client.get(self.uid)
            if history_json:
                history = json.loads(history_json)
                # 过滤掉助手回复为空的对话组，并移除时间戳
                filtered_history = []
                for item in history[-self.history_len:]:
                    # 检查是否为有效的对话组（长度至少为2）并且助手的回复（第二个元素）不为空
                    if len(item) >= 2 and item[1]['role'] == 'assistant' and item[1].get("content") != "":
                        filtered_history.append(item[:2])
                if filtered_history:
                    return numpy.concatenate(filtered_history).tolist()
                else:
                    return []
            else:
                return []
        except Exception as e:
            logger.error(f"Failed to get history: {e}")
            self.clear_history()
            return []

    def add_to_history(self, message: List):
        """
        Adds a new message to the user's chat history in Redis.

        The method retrieves the current chat history, appends the new message,
        and saves the updated history back to Redis.

        Parameters
        ----------
        message : List
            a list representing the new message to be added to the history
        """
        try:
            # 尝试获取当前历史记录
            history_json = self.redis_client.get(self.uid)
            if history_json:
                history = json.loads(history_json)
            else:
                history = []
            # 添加新的历史对话记录
            history.append(message)
            # 保存更新后的历史记录到Redis
            self.redis_client.set(self.uid, json.dumps(history))
        except Exception as e:
            logger.error(f"Failed to add to history: {e}")
            self.clear_history()

    def clear_history(self):
        """
        Clears the user's chat history from Redis.

        The method deletes the user's chat history from Redis.
        """
        try:
            self.redis_client.delete(self.uid)
        except Exception as e:
            logger.error(f"Failed to clear history: {e}")
